#pragma once

#include <cstdint>
#include <memory>
#include <mpi.h>
#include <string>
#include <vector>

namespace compi
{

// 前向声明
class MessageTagAllocator;
class TagLease;

/**
 * 标签范围描述符
 * 每个命名空间拥有一个或多个不重叠的标签范围
 */
struct TagRange
{
    int start; // 起始标签值（包含）
    int end;   // 结束标签值（不包含）

    TagRange(int start_val, int end_val) : start(start_val), end(end_val) {}

    /**
     * 检查标签值是否在此范围内
     * @param tag_value 要检查的标签值
     * @return 如果标签在范围内返回true
     */
    bool contains(int tag_value) const noexcept { return tag_value >= start && tag_value < end; }

    /**
     * 获取范围大小
     * @return 范围内可用标签数量
     */
    int size() const noexcept { return end - start; }
};

/**
 * 改进的消息标签
 * 包含标签值和所属命名空间ID
 */
class MessageTag
{
public:
    int value;           // MPI标签值
    std::uint32_t ns_id; // 命名空间ID

    MessageTag(int tag_value, std::uint32_t namespace_id) : value(tag_value), ns_id(namespace_id) {}

    /**
     * 隐式转换为int，用于MPI调用
     */
    operator int() const noexcept { return value; }

    bool operator==(const MessageTag & other) const noexcept
    {
        return value == other.value && ns_id == other.ns_id;
    }

    bool operator!=(const MessageTag & other) const noexcept { return !(*this == other); }
};

/**
 * 标签命名空间
 * 支持动态范围扩展和真正的隔离。
 *
 * 创建者与拥有者：
 * - 仅能通过 `MessageTagAllocator::create_namespace(...)` 创建；
 * - 销毁命名空间应调用 `MessageTagAllocator::destroy_namespace(ns_id)`；
 * - `TagNamespace` 自身是一个轻量句柄，内部持有指向创建它的分配器的指针，
 *   不负责命名空间资源的真正销毁与回收；生命周期归属完全由分配器统一管理。
 *
 * 统一入口：
 * - 分配器获取推荐使用 `Context::for_comm(comm).tag_allocator()`，避免直接调用
 *   `MessageTagAllocator::for_comm(comm)`（已标记为 deprecated，仅兼容旧代码）。
 *
 * 与 TransportRequest 的完成策略联动：
 * - 非阻塞原语（如 `isend/irecv`）在请求内部会对标签执行 `retain`；
 * - 请求完成时，若启用了 `CompletionStrategy::StrategyRelease`（默认开启），
 *   将自动执行 `release`（优先通过命名空间句柄，否则回退到分配器接口）；
 * - 可通过 `TransportRequest::set_strategy(...)` 清除 `StrategyRelease` 位，
 *   以禁用自动释放并由调用方自行管理标签生命周期。
 */
class TagNamespace
{
public:
    /**
     * 获取命名空间ID
     * @return 命名空间的唯一标识符
     */
    std::uint32_t id() const noexcept;

    /**
     * 获取命名空间名称
     * @return 命名空间的可选名称
     */
    const std::string & name() const noexcept;

    /**
     * 请求一个新的消息标签
     * @return 分配的消息标签
     * @throws std::runtime_error 如果无法分配标签
     */
    MessageTag request() const;

    /**
     * 增加标签的引用计数
     * @param tag 要增加引用的标签
     */
    void retain(const MessageTag & tag) const;

    /**
     * 减少标签的引用计数，可能释放标签
     * @param tag 要释放的标签
     */
    void release(const MessageTag & tag) const;

    /**
     * 检查标签是否属于此命名空间
     * @param tag_value MPI标签值
     * @return 如果标签属于此命名空间返回true
     */
    bool owns(int tag_value) const;

    /**
     * 获取命名空间的标签范围列表
     * @return 此命名空间拥有的所有标签范围
     */
    std::vector<TagRange> get_ranges() const;

    /**
     * 获取命名空间统计信息
     * @return 包含已分配、可用、总计标签数的结构
     */
    struct Stats
    {
        int allocated_count; // 已分配标签数
        int available_count; // 可用标签数
        int total_capacity;  // 总容量
        int range_count;     // 范围数量
    };
    Stats get_stats() const;

    /**
     * 非阻塞探测是否有来自指定源的消息
     * @param comm MPI通信器
     * @param src_rank 源进程排名
     * @param out_tag 输出参数，接收到的消息标签
     * @return 如果有消息返回true
     */
    bool iprobe(MPI_Comm comm, int src_rank, MessageTag & out_tag) const;

    /**
     * 申请一个一次性标签租约（RAII）
     * 返回的 TagLease 在析构或调用 retire() 时自动释放标签；
     * 可通过 renew() 重新申请一个同命名空间的新标签。
     */
    TagLease lease() const;

private:
    friend class MessageTagAllocator;

    std::uint32_t _ns_id;
    // 使用 shared_ptr 持有分配器但不管理其生命周期（自定义空删除器）
    std::shared_ptr<MessageTagAllocator> _owner;

    TagNamespace(std::uint32_t ns_id, MessageTagAllocator * allocator);
};

/**
 * TagLease：单次使用的标签句柄（RAII）
 * - 析构时若仍“活跃”，自动调用 release 归还标签
 * - retire() 主动归还并使租约失效
 * - renew() 在同命名空间重新申请一个新标签并变为“活跃”
 * - 可移动不可拷贝；移动后源对象失效
 */
class TagLease
{
public:
    TagLease() = default;
    ~TagLease();

    // 移动语义：转移租约所有权
    TagLease(TagLease && other) noexcept;
    TagLease & operator=(TagLease && other) noexcept;

    TagLease(const TagLease &) = delete;
    TagLease & operator=(const TagLease &) = delete;

    // 当前是否持有有效标签
    bool active() const noexcept { return _active; }

    // 访问当前标签
    const MessageTag & tag() const noexcept { return _tag; }
    int value() const noexcept { return _tag.value; }
    std::uint32_t ns_id() const noexcept { return _ns_id; }

    // 主动归还标签并失效（幂等）
    void retire() noexcept;

    // 在同命名空间重新申请标签；若此前活跃会先归还旧标签
    // 返回是否成功申请
    bool renew();

private:
    friend class TagNamespace;
    friend class MessageTagAllocator;

    std::shared_ptr<MessageTagAllocator> _owner{}; // 非拥有型 shared_ptr（空删除器）
    std::uint32_t _ns_id{0};
    MessageTag _tag{0, 0};
    bool _active{false};

    TagLease(std::shared_ptr<MessageTagAllocator> owner,
             std::uint32_t ns_id,
             const MessageTag & tag,
             bool active) noexcept
        : _owner(std::move(owner)), _ns_id(ns_id), _tag(tag), _active(active)
    {
    }
};

/**
 * 消息标签分配器
 * 实现真正的分层标签分配和命名空间隔离
 *
 * 设计原则：
 * 1. 每个命名空间拥有独立的标签范围，完全隔离
 * 2. 支持动态范围扩展，无固定限制
 * 3. 高效的标签分配和回收机制
 * 4. 线程安全的并发访问
 */
class MessageTagAllocator
{
    friend class TagNamespace;
    friend class ResourceManager;

public:
    /**
     * 创建新的标签命名空间
     * @param name 可选的命名空间名称
     * @return 新创建的命名空间
     */
    TagNamespace create_namespace(const std::string & name = "");

    /**
     * 根据名称查找命名空间
     * @param name 命名空间名称
     * @return 找到的命名空间，如果不存在则抛出异常
     * @throws std::runtime_error 如果命名空间不存在
     */
    TagNamespace find_namespace(const std::string & name);

    /**
     * 销毁命名空间并释放其所有资源
     * @param ns_id 要销毁的命名空间ID
     *
     * 说明：命名空间的资源由分配器统一持有和管理，调用此方法将彻底释放
     * 与该命名空间相关的范围、分配状态与引用计数。`TagNamespace` 句柄本身
     * 不执行销毁逻辑；请在确保没有悬挂引用后再调用。
     */
    void destroy_namespace(std::uint32_t ns_id);

    /**
     * 获取MPI标签上界
     * @return MPI_TAG_UB的值
     */
    int tag_ub() const noexcept;

    /**
     * 获取分配器统计信息
     */
    struct GlobalStats
    {
        int total_namespaces;     // 总命名空间数
        int total_allocated_tags; // 总已分配标签数
        int total_available_tags; // 总可用标签数
        int tag_ub_value;         // MPI标签上界
    };
    GlobalStats get_global_stats() const;

    /**
     * 请求为指定命名空间分配标签
     * @param ns_id 命名空间ID
     * @return 分配的消息标签
     * @throws std::runtime_error 如果无法分配标签
     */
    MessageTag request(std::uint32_t ns_id);

    /**
     * 增加标签的引用计数
     * @param tag 要增加引用的标签
     */
    void retain(const MessageTag & tag);

    /**
     * 减少标签的引用计数，可能释放标签
     * @param tag 要释放的标签
     */
    void release(const MessageTag & tag);

    /**
     * 检查标签是否属于指定命名空间
     * @param ns_id 命名空间ID
     * @param tag_value MPI标签值
     * @return 如果标签属于指定命名空间返回true
     */
    bool belongs(std::uint32_t ns_id, int tag_value) const;

private:
    struct Impl;
    std::unique_ptr<Impl> _impl;

    explicit MessageTagAllocator(MPI_Comm comm);
    ~MessageTagAllocator();

    // 禁止拷贝和移动
    MessageTagAllocator(const MessageTagAllocator &) = delete;
    MessageTagAllocator & operator=(const MessageTagAllocator &) = delete;
    MessageTagAllocator(MessageTagAllocator &&) = delete;
    MessageTagAllocator & operator=(MessageTagAllocator &&) = delete;
};
} // namespace compi